#!/usr/bin/python
# Needs 2 files
# 1st - tab file
#
# 2nd - COG categories file
#

import sys
import re

tab_file = sys.argv[1]
tf = open(tab_file, "rU")
lines = tf.readlines()

cog_file = sys.argv[2]
cf = open(cog_file, "rU")
clines = cf.readlines()

m1 = re.compile('\[(.*)\]\t(\w+)\t(.*)')
cogs = {}
for line in clines:
    pattern = m1.match(line)
    cogs[pattern.group(2)] = pattern.group(1)

print "seqname" + "\t" + "source" + "\t" + "feature" + "\t" + "start" + \
    "\t" + "end" + "\t" + "score" + "\t" + "strand" + "\t" + "frame"

for line in lines:
    l = line.split('\t')
    feat = l[0]
    locus_tag = l[1]
    start = l[2]
    stop = l[3]
    frame = l[4]
    gene = l[5]
    protein_id = l[6]
    ec_num = l[7]
    cog = l[8]
    bound_moiety = l[9]
    desc = l[10].rstrip()
    cog_cat = "X"

    if feat == "CDS":
        if cog != "-":
            m = re.compile('(\w+)\s.*')
            pattern = m.match(cog)
            cog_id = pattern.group(1)
            if cog_id in cogs:
                cog_cat = cogs[cog_id]
            else:
                cog_cat = "X"
            print locus_tag + "\t" + "" + "\t" + "" + cog_cat[0] + "\t" + start + \
                "\t" + stop + "\t" + "" + frame
        else:
            print locus_tag + "\t" + "" + "\t" + "" + cog_cat[0] + "\t" + start + \
                "\t" + stop + "\t" + "" + frame
    elif feat == "tRNA":
        print locus_tag + "\t" + "" + "\t" + "" + "tRNA" + "\t" + start + \
            "\t" + stop + "\t" + "" + frame
    elif feat == "rRNA":
        print locus_tag + "\t" + "" + "\t" + "" + "rRNA" + "\t" + start + \
            "\t" + stop + "\t" + "" + frame
tf.close()